# main.py (modified)
# 注意：需要安装 opencv-contrib-python 以支持 cv2.face 模块
# pip install opencv-contrib-python --upgrade
# 同时需要安装 PyQt5
# pip install pyqt5

import json
import hqyj_mqtt
import queue
import base64
import numpy as np
import cv2
import time
import traceback
import threading
from pid import PID
from tools import *  # 导入 tools 模块，包含 yolo_detect
from line_Detect import *
import sys
from PyQt5.QtWidgets import QApplication, QMessageBox
import os


"""将 MQTT 消息中携带的 Base64 编码图像数据，转换为 OpenCV 能直接处理的 BGR 格式图像矩阵（numpy.ndarray）"""
def b64_to_np(image):
    if not isinstance(image.get('image'), str):
        raise ValueError("Invalid or missing 'image' key in MQTT data")
    img_data = base64.b64decode(image['image'])        #提取Base64数据 并使用Base64解码
    img_np = np.frombuffer(img_data, dtype=np.uint8)   #将字节流转换numpy数组
    img = cv2.imdecode(img_np, cv2.IMREAD_COLOR)       #使用cv2.imdecode将NumPy数组解码为OpenCV的图像矩阵（BGR格式）
    if img is None:
        raise ValueError("Failed to decode image from Base64 data")
    return img

"""
自动控制函数，接收图像矩阵和 MQTT 客户端对象，返回当前车辆位置和图像中心位置
处理自动行驶逻辑，包括车道中心计算、PID控制、MQTT消息发送
"""
def auto_run(image, mqtt_client, pts_middle, pid, carspeed=20):
    if pts_middle is None:
        print("pts_middle is None, skipping control")
        return None, None
    lane_center = pts_middle[240:, 0].mean() #首先筛选出图像下半部分（y坐标>=240）的点，然后取这些点的x坐标，计算平均值作为车道中心
    image_center = image.shape[1] // 2   #得到图像的水平中心位置
    print(f"Frame {frame}: lane_center={lane_center}, image_center={image_center}, pts_middle shape={pts_middle.shape}")
    if np.isnan(lane_center):
        print("Invalid lane_center, skipping control")
        return lane_center, image_center

    # 使用PID控制器计算转向角，通过车道中心位置偏差进行控制
    # 输入：lane_center - 车道中心位置偏差值
    # 输出：steering_angle - 计算得到的转向角控制量
    # 说明：负号表示转向方向与偏差方向相反，实现负反馈控制
    steering_angle = -pid(lane_center) # PID控制计算转向角


    print(f"Calculated steering_angle={steering_angle}") # 打印计算出的转向角。
    mqtt_client.send_mqtt(json.dumps({"carSpeed": carspeed})) #通过MQTT客户端发送车辆速度和转向角控制指令。
    mqtt_client.send_mqtt(json.dumps({"carDirection": steering_angle}))


    return lane_center, image_center


"""
主函数，处理 MQTT 消息，并调用自动控制函数
处理MQTT重连
"""
def mqtt_reconnect(mqtt_client, host, port, username, password, q_mqtt_data):
    print("Attempting to reconnect to MQTT server...")
    try:
        mqtt_client = hqyj_mqtt.MQTTClient(host, port, username, password, q_mqtt_data)
        mqtt_client.connect()
        return mqtt_client
    except Exception as e:
        print(f"MQTT reconnect failed: {e}")
        return None

"""
人脸识别验证
"""
import time


def verify_face():

    # 加载Haar级联分类器用于人脸检测
    # 该分类器基于预训练的XML文件，专门用于检测正面人脸
    # cv2.data.haarcascades提供OpenCV内置的级联分类器文件路径
    # 'haarcascade_frontalface_default.xml'是检测正面人脸的默认模型文件
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

    # 创建LBPH人脸识别器实例
    # 该函数用于初始化一个局部二值模式直方图(LBPH)人脸识别器
    # LBPH是一种基于纹理特征的人脸识别算法，具有较好的光照不变性
    # 返回值: cv2.face_LBPHFaceRecognizer对象，用于后续的人脸训练和识别操作
    recognizer = cv2.face.LBPHFaceRecognizer_create()

    faces = []
    labels = []
    label_id = 0
    id_name_dict = {}

    if not os.path.exists('faces') or len(os.listdir('faces')) == 0:
        return False, "No enrolled faces found. Please enroll faces first."

    # 加载人脸数据
    for person_name in os.listdir('faces'):    # 遍历faces目录下的所有子目录
        person_dir = os.path.join('faces', person_name) # 构建子目录路径
        if not os.path.isdir(person_dir): # 如果不是目录，则跳过
            continue
        id_name_dict[label_id] = person_name # 创建一个字典，将标签ID映射为对应的人名
        for img_name in os.listdir(person_dir): # 遍历人脸图片
            img_path = os.path.join(person_dir, img_name) # 构建图片路径
            img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) # 读取图片数据
            if img is not None:
                faces.append(img)  # 添加图片数据
                labels.append(label_id)   # 为每个图片分配一个唯一的标签ID
        label_id += 1

    if not faces:
        return False, "No valid face images loaded."

    # 训练识别器
    recognizer.train(faces, np.array(labels)) # 使用训练数据训练识别器

    # 开始识别
    cap = cv2.VideoCapture(0) # 打开摄像头
    verified = False # 默认为未验证
    frame_count = 0 # 帧计数器
    max_frames = 300  # 大约10秒，假设30fps

    start_time = time.time()  # 记录开始时间
    min_duration = 3.0  # 最小持续时间(秒)

    while frame_count < max_frames: # 循环处理每一帧
        ret, frame = cap.read() # 读取一帧
        if not ret: # 如果没有读取到帧，则退出循环
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # 转换为灰度图像
        detected_faces = face_cascade.detectMultiScale(gray, 1.3, 5) # 检测人脸

        for (x, y, w, h) in detected_faces: # 遍历检测到的人脸
            face_roi = gray[y:y + h, x:x + w] # 获取人脸区域
            label, confidence = recognizer.predict(face_roi) # 预测人脸
            if confidence < 100: # 如果置信度小于100，则认为人脸已验证  低置信度表示高相似度
                person_name = id_name_dict.get(label, "Unknown") # 获取人名
                verified = True # 设置已验证状态
                cv2.putText(frame, f'Verified: {person_name} ({confidence:.2f})', (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, # 在人脸区域绘制文字
                            0.9, (0, 255, 0), 2) # 设置文字颜色和字体大小
            else: # 否则，认为人脸未验证
                cv2.putText(frame, f'Unknown ({confidence:.2f})', (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, # 在人脸区域绘制文字
                            (0, 0, 255), 2) # 设置文字颜色和字体大小
            cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2) # 绘制人脸区域

        cv2.imshow('Face Verification', frame)

        # 检查是否满足最小时间要求
        elapsed_time = time.time() - start_time
        if (cv2.waitKey(1) & 0xFF == ord('q')) or (verified and elapsed_time >= min_duration):
            break

        frame_count += 1

    cap.release()
    cv2.destroyAllWindows()
    return verified, "人脸识别成功，车辆启动中..." if verified else "人脸识别失败，无法启动车辆！"


if __name__ == '__main__':
    # 使用 PyQt5 创建前端界面
    app = QApplication(sys.argv)

    verified, message = verify_face() # 人脸识别

    if not verified:
        QMessageBox.critical(None, "错误", message)
        sys.exit()

    QMessageBox.information(None, "成功", message)

    q_mqtt_data = queue.Queue(10)  # 增加队列容量
    frame = 0
    mqtt_client = hqyj_mqtt.MQTTClient('127.0.0.1', 21883, 'bb', 'aa', q_mqtt_data)  # 初始化MQTT客户端
    # 初始化PID控制器，调整 PID 参数
    pid = PID(Kp=0.25, Ki=0.0, Kd=0.0, setpoint=240) # setpoint：PID尝试达到的初始设定点   高速场景需更强的Kd微分项抑制振荡，低速场景需更大的积分项Ki消除静差
    pid.sample_time = 0.1  # 控制器生成新输出值之前应等待的时间
    pid.output_limits = (-13, 13)

    while True:
        try:
            image_data = q_mqtt_data.get(timeout=1.0)
            if 'image' not in image_data:
                print("Invalid MQTT data, skipping")
                continue
            image = b64_to_np(image_data)  # 将MQTT消息中携带的Base64编码图像数据，转换为OpenCV能直接处理的BGR格式图像矩阵
            print(f"Image loaded: shape={image.shape}")
            image_warp, minv = perspective_transform(image)  # 透视变换 获取鸟瞰图image_warp
            if image_warp is None:
                print("Perspective transform failed, skipping frame")
                continue
            # 图像增强
            img_close = image_enhance(image_warp)
            # 停止线检测
            stop_line_flag = find_stop_line(image_warp.copy())
            print(f"stop_line_flag={stop_line_flag}")  # 调试信息
            # YOLO 检测
            light_color, _ = yolo_detect(image)  # 忽略 person_warning
            print(f"light_color={light_color}")  # 调试信息

            if stop_line_flag:
                mqtt_client.send_mqtt(json.dumps({"carSpeed": 0}))
                print('识别到的颜色为:', light_color)
                if light_color is not None:
                    if light_color == 0 or light_color == 1:  # 红灯或黄灯
                        print('停止车辆')
                        mqtt_client.send_mqtt(json.dumps({"carSpeed": 0}))
                        cv2.waitKey(1)
                        continue
                    elif light_color == 2:  # 绿灯
                        print('绿灯，启动车辆')
                        mqtt_client.send_mqtt(json.dumps({"carSpeed": 20}))
                        stop_line_flag = False  # 重置停止标志

            # 车道线检测
            lineLoc = find_line(img_close)
            if lineLoc is None:
                print("Lane detection failed, skipping frame")
                continue
            pts_middle = show_line(image, image_warp, lineLoc, minv)
            lane_center, image_center = auto_run(image, mqtt_client, pts_middle, pid)
            frame += 1
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        except queue.Empty:
            print("No data from MQTT queue, continuing...")
            # 尝试重连
            mqtt_client = mqtt_reconnect(mqtt_client, '127.0.0.1', 21883, 'bb', 'aa', q_mqtt_data)
            if mqtt_client is None:
                time.sleep(1)  # 等待后重试
            frame += 1
        except Exception as e:
            print(f"Error in main loop: {e}")
            traceback.print_exc()
            time.sleep(0.1)

    cv2.destroyAllWindows()

    sys.exit(app.exec_())